#include "../DWMRICore/DTI.h"
#include "../DWMRICore/Scalar.h"

#include <teem/ten.h>

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define _USE_MATH_DEFINES
#include <math.h>
#include <time.h>


char *dti_pathname = NULL;
char *mask_pathname = NULL;
char *dst_pathname = NULL;

CDTI *pDTI = NULL;
CScalar *pMASK = NULL;

float threshold = 0.2f;

/* MLS fiber tracking parameters */
float MLSSize = 2.0f;
CPoint3F *pSamplePoints = NULL;
float *pSamplePointsWeight = NULL;



void PrintUsage()
{
	printf("Usage: TwoTensorModelSelection -dti dti_file -mask mask_file -dst dst_file [-thresh threshold]\n");
}

int ParseArguments(const int argc, char *argv[])
{
	for (int i = 1; i < argc; ++i) {
		if (strcmp(argv[i], "-dti") == 0) {
			dti_pathname = argv[++i];
		} else if (strcmp(argv[i], "-dst") == 0) {
			dst_pathname = argv[++i];
		} else if (strcmp(argv[i], "-mask") == 0) {
			mask_pathname = argv[++i];
		} else if (strcmp(argv[i], "-thresh") == 0) {
			threshold = (float)atof(argv[++i]);
		} else {
			return -1;
		}
	}

	if (dti_pathname == NULL || mask_pathname == NULL || dst_pathname == NULL)
		return -2;

	return 0;
}

float UpdateP(const uint x, const uint y, const uint z, 
	const uint w, const uint h, const uint d, float *p, float *np)
{
	uint min_x = x > 0 ? x-1 : 0;
	uint max_x = x < w-1 ? x+1 : w-1;
	uint min_y = y > 0 ? y-1 : 0;
	uint max_y = y < h-1 ? y+1 : h-1;
	uint min_z = z > 0 ? z-1 : 0;
	uint max_z = z < d-1 ? z+1 : d-1;

	float dp = 0.0f;
	float dnp = 0.0f;
	uint index = 0;
	for (uint cz = min_z; cz <= max_z; ++cz) {
		for (uint cy = min_y; cy <= max_y; ++cy) {
			for (uint cx = min_x; cx <= max_x; ++cx) {
				if (pMASK->GetScalar(cx, cy, cz) > 0.5f) {
					index = (cz * h + cy) * w + cx;
					dp += p[index];
					dnp += np[index];
				}
			}
		}
	}

	index = (z * h + y) * w + x;
	float res = p[index] * (1.0f + dp) / (p[index] * (1.0f + dp) + np[index] * (1.0f + dnp));
	return res;
}

float UpdateNP(const uint x, const uint y, const uint z, 
	const uint w, const uint h, const uint d, float *p, float *np)
{
	uint min_x = x > 0 ? x-1 : 0;
	uint max_x = x < w-1 ? x+1 : w-1;
	uint min_y = y > 0 ? y-1 : 0;
	uint max_y = y < h-1 ? y+1 : h-1;
	uint min_z = z > 0 ? z-1 : 0;
	uint max_z = z < d-1 ? z+1 : d-1;

	float dp = 0.0f;
	float dnp = 0.0f;
	uint index = 0;
	for (uint cz = min_z; cz <= max_z; ++cz) {
		for (uint cy = min_y; cy <= max_y; ++cy) {
			for (uint cx = min_x; cx <= max_x; ++cx) {
				if (pMASK->GetScalar(cx, cy, cz) > 0.5f) {
					index = (cz * h + cy) * w + cx;
					dp += p[index];
					dnp += np[index];
				}
			}
		}
	}

	index = (z * h + y) * w + x;
	float res = np[index] * (1.0f + dnp) / (p[index] * (1.0f + dp) + np[index] * (1.0f + dnp));
	return res;
}
void CreateMLSSamplePoints()
{
	SafeDeleteArray(pSamplePoints);
	SafeDeleteArray(pSamplePointsWeight);
	pSamplePoints = new CPoint3F[SAMPLE_POINT_NUM];
	pSamplePointsWeight = new float[SAMPLE_POINT_NUM];

	pSamplePoints[0].m_x = 0.0f;
	pSamplePoints[0].m_y = 0.0f;
	pSamplePoints[0].m_z = 0.0f;
	pSamplePointsWeight[0] = 1.0f;

	clock_t seed = clock();
	srand(seed);

	for (int i = 1; i < SAMPLE_POINT_NUM; ++i) {
		bool flag = false;
		CPoint3F pt;
		while(!flag) {
			float a = (float)rand() / (float)RAND_MAX * 2.0f * (float)M_PI;
			float b = (float)rand() / (float)RAND_MAX * (float)M_PI;
			float r = (float)rand() / (float)RAND_MAX;

			pt.m_x = r * sinf(a) * sinf(b);
			pt.m_y = r * sinf(a) * cosf(b);
			pt.m_z = r * cosf(a);

			flag = true;
			for (int j = 0; j < i; ++j) {
				float dis = Distance(pSamplePoints[j], pt);
				if (dis < 0.25f) {
					flag = false;
					break;
				}
			}
		}

		pSamplePoints[i] = pt;
		pSamplePointsWeight[i] = InnerProduct(pSamplePoints[i], pSamplePoints[i]) * 4.0f;
		pSamplePointsWeight[i] = expf(pSamplePointsWeight[i] / -2.5f);
	}
}

void DTIMLSFilter(const CPoint3F pos, float *lastDTI, float *filteredDTI)
{
	float m[9];
	TEN_T2M(m, lastDTI);

	CPoint3F samples[SAMPLE_POINT_NUM];
	float length[SAMPLE_POINT_NUM];
	float max_length = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = pSamplePoints[i].m_x * m[0] + pSamplePoints[i].m_y * m[1] + pSamplePoints[i].m_z * m[2];
		samples[i].m_y = pSamplePoints[i].m_x * m[3] + pSamplePoints[i].m_y * m[4] + pSamplePoints[i].m_z * m[5];
		samples[i].m_z = pSamplePoints[i].m_x * m[6] + pSamplePoints[i].m_y * m[7] + pSamplePoints[i].m_z * m[8];
		length[i] = sqrtf(InnerProduct(samples[i], samples[i]));
		if (length[i] > max_length)
			max_length = length[i];
	}

	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = samples[i].m_x / max_length * MLSSize + pos.m_x;
		samples[i].m_y = samples[i].m_y / max_length * MLSSize + pos.m_y;
		samples[i].m_z = samples[i].m_z / max_length * MLSSize + pos.m_z;
	}

	memset(filteredDTI, 0, sizeof(float)*7);
	float temp_DTI[7];
	float weight = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		if (pMASK->GetScalar(samples[i]) < 0.5f) {
			continue;
		}
		pDTI->GetDTI(samples[i], temp_DTI);

		for (uint j = 0; j < 7; ++j) {
			filteredDTI[j] += pSamplePointsWeight[i] * temp_DTI[j];
		}
		weight += pSamplePointsWeight[i];
	}

	for (uint j = 0; j < 7; ++j) {
		filteredDTI[j] = filteredDTI[j] / weight;
	}
}


void ModelSelection(CDTI *pDTI, CScalar *pMASK, CScalar *res)
{
	uint w = pDTI->GetWidth();
	uint h = pDTI->GetHeight();
	uint d = pDTI->GetDepth();	
	
	float *cp = new float[w*h*d];

	// compute CP for each voxel
	int i = 0;
	for (uint z = 0; z < d; ++z) {
		for (uint y = 0; y < h; ++y) {
			for (uint x = 0; x < w; ++x) {
				CPoint3F pos = CPoint3F(x, y, z);
				float *initDTI = new float[7];
				float *filteredDTI = new float[7];
				pDTI->GetDTIFast(x, y, z, initDTI);
				DTIMLSFilter(pos, initDTI, filteredDTI);

				if (pMASK->GetScalarFast(i) > 0.5f) {
					cp[i] = pDTI->DTI2CP(filteredDTI);
				} else {
					cp[i] = 0.0f;
				}
				
				SafeDeleteArray(initDTI);
				SafeDeleteArray(filteredDTI);

				i++;
			}
		}
	}
	
	float *temp_p = new float[w*h*d];
	float *temp_np = new float[w*h*d];

	// do the model selection for cp
	printf("Model Selectionf or CP ... \n");
	float *cp_p = new float[w*h*d];
	float *cp_np = new float[w*h*d];
	for (uint i = 0; i < w*h*d; ++i) {
		cp_p[i] = 0.5f + (cp[i] - threshold) / (1.0f - threshold) * 0.5f;
		cp_np[i] = 1.0f - cp_p[i];
	}

	while (true) {
		for (uint z = 0; z < d; ++z) {
			for (uint y = 0; y < h; ++y) {
				for (uint x = 0; x < w; ++x) {
					uint index = (z * h + y) * w + x;
					temp_p[index] = UpdateP(x, y, z, w, h, d, cp_p, cp_np);
					temp_np[index] = UpdateNP(x, y, z, w, h, d, cp_p, cp_np);
					if ((temp_p[index] + temp_np[index]) < 0.999f) {
						printf("%f %f\n", temp_p[index], temp_np[index]);
					}
					//temp_np[index] = 1.0f - temp_p[index];
				}
			}
		}

		float average_delta = 0.0f;
		for (uint i = 0; i < w*h*d; ++i) {
			float delta = fabs(cp_p[i] - temp_p[i]);
			average_delta += delta;
		}
		average_delta = average_delta / (float)(w*h*d);
		printf("average delta %f\n", average_delta);

		memcpy(cp_p, temp_p, sizeof(float)*w*h*d);
		memcpy(cp_np, temp_np, sizeof(float)*w*h*d);
		
		if (average_delta < 0.0001f)
			break;
	}

	int count = 0;
	int total = 0;
	for (uint i = 0; i < w*h*d; ++i) {
		if (pMASK->GetScalarFast(i) > 0.5f) {
			if (cp_p[i] > 0.5f) {
				res->SetScalar(i, 1.0f);
				count++;
			} else {
				res->SetScalar(i, 0.0f);
			}
			total++;
		} else {
			res->SetScalar(i, 0.0f);
		}	
	}
	float ratio = (float)count / (float)total;
	printf("ratio = %f\n", ratio);

	SafeDeleteArray(cp);
	SafeDeleteArray(cp_p);
	SafeDeleteArray(cp_np);
	SafeDeleteArray(temp_p);
	SafeDeleteArray(temp_np);
}

int main(int argc, char *argv[])
{
	if (ParseArguments(argc, argv) != 0) {
		PrintUsage();
		return 0;
	}

	pDTI = new CDTI();
	pDTI->OpenDTIFile(dti_pathname);
	pMASK = new CScalar();
	pMASK->OpenScalarFile(mask_pathname);

	uint w = pDTI->GetWidth();
	uint h = pDTI->GetHeight();
	uint d = pDTI->GetDepth();

	CScalar *res = new CScalar();
	res->CreateScalar(w, h, d);

	CreateMLSSamplePoints();

	ModelSelection(pDTI, pMASK, res);

	res->SaveScalarFile(dst_pathname);

	SafeDelete(pDTI);
	SafeDelete(pMASK);
	SafeDelete(res);

	return 0;
}